import os
import pathlib
import argparse
import numpy as np
from numba import jit
import PIL.Image
from tqdm import tqdm


def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument('--image-dir', type=str, help='image directory', required=True)
    parser.add_argument('--output-dir', type=str, help='output directory', required=True)
    parser.add_argument('--nearest-map', type=str, help='path to nearest map', required=True)
    args = parser.parse_args()
    return args


@jit(nopython=True, nogil=True)
def convert(image_rgb, nearest_map):
    h, w = image_rgb.shape[:2]
    image_out = np.empty((h, w), dtype=np.uint8)
    for i in range(h):
        for j in range(w):
            image_out[i, j] = convert_pixel(image_rgb[i, j], nearest_map)
    return image_out


@jit(nopython=True, cache=True)
def convert_pixel(pixel, nearest_map):
    return nearest_map[pixel[0], pixel[1], pixel[2]]


if __name__ == '__main__':
    args = parse_args()

    nearest_map = np.load(args.nearest_map, allow_pickle=True)

    image_paths = [str(path) for path in pathlib.Path(args.image_dir).rglob('*.png')]
    for image_path in tqdm(image_paths):
        image = PIL.Image.open(image_path)
        image = np.asarray(image)
        image = convert(image, nearest_map)

        output_path = args.output_dir + image_path[len(args.image_dir):]
        os.makedirs(os.path.dirname(output_path), exist_ok=True)
        PIL.Image.fromarray(image).save(output_path)

    print("done")
